import numpy as np
from os.path import join
from torch.utils.data import Dataset
from PIL import Image
import data.vfitransforms as vt
import os
import cv2


class Haworth(Dataset):
    def __init__(
        self,
        db_dir,
    ):
        self.db_dir = db_dir
        self.data = self._prepare(db_dir)

    def __len__(self):
        return len(self.data)

    def _prepare(self, db_dir):
        data = []
        root_dir = join(db_dir, "images")
        for file in os.listdir(root_dir):
            if file.endswith(".tif"):
                data.append((join("images", file), -1))
        return data

    def _reader(self, path):
        cat = cv2.imread(join(self.db_dir, path), cv2.IMREAD_UNCHANGED)
        return cat

    def __getitem__(self, index):
        raw_path, label = self.data[index]
        return {"raw": self._reader(raw_path), "cls": label}


class MDCD(Dataset):
    def __init__(
        self,
        db_dir,
        channels,
        augment_s=True,
        augment_t=True,
    ):
        self.channels = channels
        self.augment_s = augment_s
        self.augment_t = augment_t
        self.db_dir = db_dir
        self.data = self._prepare(db_dir)

    def __len__(self):
        return len(self.data)

    def _prepare(self, db_dir):
        raise NotImplementedError

    def _reader(self, path):
        cat = Image.open(join(self.db_dir, path))
        if self.augment_s:
            cat = vt.rand_flip(cat, p=0.5)
        if self.augment_t:
            cat = vt.rand_reverse(cat, p=0.5)
        cat = np.array(cat, dtype=np.float32).squeeze()
        return cat / 127.5 - 1.0

    def __getitem__(self, index):
        raw_path, label = self.data[index]
        return {"raw": self._reader(raw_path), "cls": label}


class MDCDTrain(MDCD):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "train.txt")) as f:
            for raw, label in map(lambda x: (x.strip(), int("non" in x)), f):
                data.append((raw, label))
        return data


class MDCDValidate(MDCD):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "val.txt")) as f:
            for raw, label in map(lambda x: (x.strip(), int("non" in x)), f):
                data.append((raw, label))
        return data


class MDCDTest(MDCD):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "val.txt")) as f:
            for raw, label in map(lambda x: (x.strip(), int("non" in x)), f):
                data.append((raw, label))
        with open(join(db_dir, "train.txt")) as f:
            for raw, label in map(lambda x: (x.strip(), int("non" in x)), f):
                data.append((raw, label))
        return data


class MDCDTestMars(MDCD):

    def _prepare(self, db_dir):
        data = []
        root_dir = join(db_dir, "images")
        for file in os.listdir(root_dir):
            if file.endswith(".jpg"):
                data.append((join("images", file), -1))
        return data


if __name__ == "__main__":
    dataset = MDCDValidate(
        db_dir="/disk527/sdb1/a804_cbf/datasets/MDCD",
        channels=3,
        augment_s=False,
        augment_t=False,
    )
    for i, data in enumerate(dataset):
        print(data["raw"].shape)
        print(data["cls"])
        if i > 10:
            break
